import pandas as pd
import random
from collections import defaultdict


def genetic_algorithm(weight_coefficients, performance_constraints, price_requirement, part_requirements,
                      accessory_requirements, data, population_size=100, generations=200, mutation_rate=0.1):
    # 如果没有提供性能权重，自动填充
    if not weight_coefficients:
        num_constraints = len(performance_constraints)
        if num_constraints > 0:
            equal_weight = 1 / num_constraints
            weight_coefficients = {performance: equal_weight for performance in performance_constraints.keys()}

    # 将数据框转换为字典，方便快速访问
    accessory_info = {}
    for _, row in data.iterrows():
        accessory = row['配件']
        accessory_info[accessory] = {
            '价格': row['价格'],
            '性能': {
                performance: row[performance]
                for performance in performance_constraints.keys()
                if pd.notnull(row[performance])
            }
        }

    # 获取所有部位
    parts = data['部位'].unique()
    part_available_accessories = {
        part: data[data['部位'] == part]['配件'].tolist()
        for part in parts
    }

    # 创建初始种群
    def create_individual():
        individual = {}
        for part in parts:
            available_parts = part_available_accessories[part]
            if available_parts:
                if part in part_requirements and part_requirements[part] =='must':
                    individual[part] = random.choice(available_parts)
                else:
                    individual[part] = random.choice(available_parts + [None])
            else:
                individual[part] = None
        for accessory, requirement in accessory_requirements.items():
            if isinstance(requirement, str) and requirement not in ['must','must_not']:
                for part in parts:
                    if data[data['配件'] == requirement]['部位'].values[0] == part:
                        individual[part] = requirement
        return individual

    population = [create_individual() for _ in range(population_size)]

    # 评估函数
    def evaluate(individual):
        total_price = 0
        performance_values = defaultdict(int)

        for accessory in individual.values():
            if accessory:
                info = accessory_info[accessory]
                total_price += info['价格']
                for performance, value in info['性能'].items():
                    performance_values[performance] += value

        # 检查价格要求
        if price_requirement[0] <= total_price <= price_requirement[1]:
            price_score = 1
        else:
            price_score = 0

        # 检查性能约束
        performance_score = 1
        for performance, constraint in performance_constraints.items():
            if performance_values[performance] < constraint:
                performance_score = 0
                break

        # 检查配件要求
        accessory_score = 1
        for accessory, requirement in accessory_requirements.items():
            if requirement =='must':
                if accessory not in individual.values():
                    accessory_score = 0
                    break
            elif requirement =='must_not':
                if accessory in individual.values():
                    accessory_score = 0
                    break
        # 检查部位要求
        for part, requirement in part_requirements.items():
            if requirement =='must':
                if individual[part] is None:
                    accessory_score = 0
                    break

        # 检查增高座瞄具和瞄具的关系
        if '增高座瞄具' in individual.values() and '瞄具' not in individual.values():
            accessory_score = 0

        # 检查枪管和贴片的关系
        if '枪管' not in individual.values():
            for patch in ['左贴片', '右贴片', '上贴片']:
                if patch in individual.values():
                    accessory_score = 0
                    break

        # 计算总得分
        if price_score == 0 or performance_score == 0 or accessory_score == 0:
            fitness = 0
        else:
            weighted_performance = sum(
                weight_coefficients[performance] * performance_values[performance]
                for performance in weight_coefficients.keys()
                if performance in performance_values
            )
            fitness = weighted_performance

        return fitness

    # 选择函数（轮盘赌选择）
    def select(population):
        fitness_scores = [evaluate(individual) for individual in population]
        total_fitness = sum(fitness_scores)
        if total_fitness == 0:
            return random.choices(population, k=2)
        probabilities = [score / total_fitness for score in fitness_scores]
        return random.choices(population, weights=probabilities, k=2)

    # 交叉函数
    def crossover(parent1, parent2):
        child = {}
        parts = list(parent1.keys())
        crossover_point = random.randint(1, len(parts) - 1)
        for i in range(crossover_point):
            child[parts[i]] = parent1[parts[i]]
        for i in range(crossover_point, len(parts)):
            child[parts[i]] = parent2[parts[i]]
        return child

    # 变异函数
    def mutate(individual):
        for part in individual.keys():
            if random.random() < mutation_rate:
                available_parts = part_available_accessories[part]
                if available_parts:
                    if part in part_requirements and part_requirements[part] =='must':
                        individual[part] = random.choice(available_parts)
                    else:
                        individual[part] = random.choice(available_parts + [None])
        return individual

    # 遗传算法主循环
    for i in range(generations):
        print(i)
        new_population = []
        for _ in range(population_size // 2):
            parent1, parent2 = select(population)
            child1 = crossover(parent1, parent2)
            child2 = crossover(parent2, parent1)
            child1 = mutate(child1)
            child2 = mutate(child2)
            new_population.extend([child1, child2])
        population = new_population

    best_individual = max(population, key=evaluate)
    best_fitness = evaluate(best_individual)

    if best_fitness == 0:
        return None
    else:
        return best_individual
    

# 各项性能权重，可为0
weight_coefficients = {
    # '优势射程': 0.1,
    # '后坐力控制': 0.4,
    # '操控速度': 0.2,
    # '据枪稳定性': 0.2,
    # '腰际射击精度': 0.1
}
# 性能要求，可不填
performance_constraints = {
    '后坐力控制': 10
}
# 价格要求，可不填
price_requirement = (100000, 200000)
# 部位要求，可不填
part_requirements = {
    '枪口':'must'
}
# 配件要求，可不填
accessory_requirements = {
    # '薪火螺旋消焰器':'must',
    # '先进多口径消音器':'must_not'
}

# 获取数据
excel_file = pd.ExcelFile(r"C:\Users\Administrator\Documents\改枪.xlsx")
data = pd.read_excel(excel_file)

# 运行遗传算法
result = genetic_algorithm(weight_coefficients, performance_constraints, price_requirement, part_requirements, accessory_requirements, data)

if result:
    print("最优配件方案:")
    for part, accessory in result.items():
        print(f"{part}: {accessory}")
else:
    print("不存在符合要求的配件方案")